{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SFTTrainer를 활용한 지도 학습 기반 미세 조정\n",
    "\n",
    "이 노트북은 `trl` 라이브러리에서 제공하는 `SFTTrainer`를 이용해 `HuggingFaceTB/SmolLM2-135M` 모델을 미세 조정하는 법을 설명합니다. 노트북 셀을 실행하면 모델이 미세 조정됩니다. 다른 데이터셋을 사용해 난이도를 조절할 수 있습니다.\n",
    "\n",
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'>\n",
    "    <h2 style='margin: 0;color:blue'>연습: SFTTrainer를 활용한 SmolLM2 미세 조정</h2>\n",
    "    <p>Hugging Face 허브에서 가져온 데이터셋으로 모델은 미세 조정해보세요. </p> \n",
    "    <p><b>난이도</b></p>\n",
    "    <p>🐢 `HuggingFaceTB/smoltalk` 데이터셋 사용해보기</p>\n",
    "    <p>🐕 `bigcode/the-stack-smol` 데이터셋의 하위 집합인 `data/python`을 활용해 코드 생성 모델 미세 조정해보기</p>\n",
    "    <p>🦁 실제 사용 사례와 관련된 관심 있는 데이터셋 사용해보기</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Google Colab에서 requirements 설치\n",
    "# !pip install transformers datasets trl huggingface_hub\n",
    "\n",
    "# Hugging Face 인증\n",
    "\n",
    "from huggingface_hub import login\n",
    "login()\n",
    "\n",
    "# 허브 토큰을 HF_TOKEN 환경 변수로 설정해두면 편하게 사용할 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 주요 라이브러리 불러오기\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from datasets import load_dataset\n",
    "from trl import SFTConfig, SFTTrainer, setup_chat_format\n",
    "import torch\n",
    "\n",
    "device = (\n",
    "    \"cuda\"\n",
    "    if torch.cuda.is_available()\n",
    "    else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
    ")\n",
    "\n",
    "# 모델과 토크나이저 불러오기\n",
    "model_name = \"HuggingFaceTB/SmolLM2-135M\"\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    pretrained_model_name_or_path=model_name\n",
    ").to(device)\n",
    "tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_name)\n",
    "\n",
    "# 대화 형식 설정\n",
    "model, tokenizer = setup_chat_format(model=model, tokenizer=tokenizer)\n",
    "\n",
    "# 미세 조정 결과를 저장하고 업로드하기 위한 이름 설정\n",
    "finetune_name = \"SmolLM2-FT-MyDataset\"\n",
    "finetune_tags = [\"smol-course\", \"module_1\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 기본 모델을 활용한 답변 생성\n",
    "\n",
    "여기서는 대화 템플릿이 없는 기본 모델을 사용해보겠습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 학습 전 기본 모델 테스트 진행\n",
    "prompt = \"Write a haiku about programming\"\n",
    "\n",
    "# 템플릿으로 메시지 형식 지정\n",
    "messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "\n",
    "# 응답 생성\n",
    "inputs = tokenizer(formatted_prompt, return_tensors=\"pt\").to(device)\n",
    "outputs = model.generate(**inputs, max_new_tokens=100)\n",
    "print(\"Before training:\")\n",
    "print(tokenizer.decode(outputs[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 데이터셋 준비\n",
    "\n",
    "샘플 데이터셋을 불러와서 학습을 위한 형식을 지정합니다. 데이터셋은 입력-출력 쌍으로 구성되어야 하며, 각 입력은 프롬프트이고 출력은 모델에서 예상되는 응답입니다. **TRL은 모델의 대화 템플릿에 맞게 입력 메시지 형식을 맞춥니다.** 메시지는 `role`과 `content` 키를 가진 딕셔너리 리스트로 표현되어야 합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 샘플 데이터셋 불러오기\n",
    "from datasets import load_dataset\n",
    "\n",
    "# TODO: path와 name 파라미터를 이용해 데이터셋과 configuration 정의하기\n",
    "ds = load_dataset(path=\"HuggingFaceTB/smoltalk\", name=\"everyday-conversations\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: 🦁 데이터셋이 TRL에서 대화 템플릿으로 변환할 수 있는 형식이 아니라면, 해당 데이터를 처리해야 합니다. 자세한 내용은 [모듈](../chat_templates.md)을 참고하세요."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SFTTrainer 설정\n",
    "\n",
    "`SFTTrainer`는 학습 단계 수, 배치 크기, 학습률, 평가 방식과 같이 학습 과정을 제어하는 다양한 파라미터로 구성됩니다. 특정 요구 사항과 연산 자원에 맞춰 파라미터를 조정하세요. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SFTTrainer를 위한 SFT configuration 설정\n",
    "sft_config = SFTConfig(\n",
    "    output_dir=\"./sft_output\",\n",
    "    max_steps=1000,  # 원하는 학습 시간과 데이터셋 크기에 따라 조정\n",
    "    per_device_train_batch_size=4,  # GPU 메모리 용량에 따라 설정\n",
    "    learning_rate=5e-5,  # 미세 조정을 위해 일반적으로 쓰이는 값\n",
    "    logging_steps=10,  # 학습 지표 로깅 빈도\n",
    "    save_steps=100,  # 모델 체크포인트 저장 빈도\n",
    "    evaluation_strategy=\"steps\",  # 주기적인 모델 평가 설정\n",
    "    eval_steps=50,  # 평가 빈도\n",
    "    use_mps_device=(\n",
    "        True if device == \"mps\" else False\n",
    "    ),  # 혼합 정밀도 학습(mixed precision training)을 위한 MPS 사용\n",
    "    hub_model_id=finetune_name,  # 모델 이름 설정\n",
    ")\n",
    "\n",
    "# SFTTrainer 초기화\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    args=sft_config,\n",
    "    train_dataset=ds[\"train\"],\n",
    "    eval_dataset=ds[\"test\"],\n",
    ")\n",
    "\n",
    "# TODO: 🦁 🐕 선택한 데이터셋에 맞게 SFTTrainer 파라미터를 조절하세요. 예를 들어, `bigcode/the-stack-smol` 데이터셋을 사용하는 경우 `content` 열을 선택해야 합니다."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 모델 학습\n",
    "\n",
    "트레이너가 구성되었기 때문에 이제 모델을 학습시킬 수 있습니다. 학습 과정은 데이터셋을 반복적으로 처리하며 손실을 계산하고, 이 손실을 최소화하기 위해 모델의 파라미터를 업데이트하는 과정을 포함합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 모델 학습\n",
    "trainer.train()\n",
    "\n",
    "# 모델 저장\n",
    "trainer.save_model(f\"./{finetune_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.push_to_hub(tags=finetune_tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'>\n",
    "    <h2 style='margin: 0;color:blue'>추가 연습: 미세 조정 모델로 응답 생성하기</h2>\n",
    "    <p>🐕 기본 예제에서와 마찬가지로 미세 조정된 모델을 사용하여 응답을 생성해보세요.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 동일한 프롬프트로 미세 조정된 모델을 테스트해보세요.\n",
    "\n",
    "# 학습 전 기본 모델 테스트 진행\n",
    "prompt = \"Write a haiku about programming\"\n",
    "\n",
    "# 템플릿으로 메시지 형식 지정\n",
    "messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "\n",
    "# 응답 생성\n",
    "inputs = tokenizer(formatted_prompt, return_tensors=\"pt\").to(device)\n",
    "\n",
    "# TODO: 기본 예제에서 했던 것처럼 미세 조정된 모델을 사용해 응답을 생성해보세요."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 💐 완료!\n",
    "\n",
    "이 노트북은 `SFTTrainer`를 사용하여 `HuggingFaceTB/SmolLM2-135M` 모델을 미세 조정하는 단계별 가이드를 제공합니다. 이 단계를 따라 특정 작업을 보다 효과적으로 수행하도록 모델을 조정할 수 있습니다. 이 과정을 계속 진행하려면 다음 단계를 시도해 보세요:\n",
    "\n",
    "- 더 어려운 방법으로 노트북 실습해보기\n",
    "- 동료의 PR 검토하기\n",
    "- 이슈 또는 PR을 통해 자료 개선하기"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
